//
//  BSDF.hpp
//  RayTracer
//
//  Created by 俞云康 on 4/5/16.
//  Copyright © 2016 yushroom. All rights reserved.
//

#ifndef BSDF_hpp
#define BSDF_hpp

#include "HitInfo.hpp"
#include "Color.hpp"
#include "Sample.hpp"
#include "BxDF.hpp"


class BSDF {
protected:
    const HitInfo mHit;
    const float mEta;
private:
    Normal mShadingNormal, mGeometryNormal;
    Direction mSn, mTn;
    int mBxDFsSize = 0;
#define MAX_BxDFS 8
    BxDF *bxdfs[MAX_BxDFS];
    
public:
    BSDF(const HitInfo& hit, const Normal &geometryNormal, const float eta = 1.f)
        : mHit(hit), mEta(eta), mGeometryNormal(geometryNormal), mTn(hit.dpdv) {
            mShadingNormal = Normalize(hit.normal);
            mSn = Normalize(hit.dpdu);
            mTn = Cross(mShadingNormal, mSn);
    }
    virtual ~BSDF() {
        for (int i = 0; i < mBxDFsSize; ++i) {
            delete bxdfs[i];
        }
    }
    
    inline void add(BxDF *bxdf) {
        Assert(mBxDFsSize < MAX_BxDFS);
        bxdfs[mBxDFsSize++] = bxdf;
    }
    
    int getComponentsCount() const { return mBxDFsSize; }
    
    int getComponentsCount(BxDFType flags) const {
        int num = 0;
        for (int i = 0; i < mBxDFsSize; ++i) {
            if (bxdfs[i]->matchesFlags(flags))
                ++num;
        }
        return num;
    }
    
    Direction world2Local(const Direction& v) const {
        auto& sn = mSn;
        auto& tn = mTn;
        auto& nn = mShadingNormal;
        return Direction(Dot(v, sn), Dot(v, tn), Dot(v, nn));
    }
    
    Direction local2World(const Direction& v) const {
        auto& sn = mSn;
        auto& tn = mTn;
        auto& nn = mShadingNormal;
        return Direction(sn.x*v.x + tn.x*v.y + nn.x*v.z,
                         sn.y*v.x + tn.y*v.y + nn.y*v.z,
                         sn.z*v.x + tn.z*v.y + nn.z*v.z);
    }
    
    virtual Color f(const Direction& woW, const Direction& wiW, BxDFType flags = BxDF_ALL) const {
        Direction wiL = world2Local(wiW);
        Direction woL = world2Local(woW);
        if (Dot(wiW, mGeometryNormal) * Dot(woW, mGeometryNormal) > 0) // ignore BTDFs
            flags = BxDFType(flags & ~BxDF_TRANSMISSION);
        else // ignore BRDFs
            flags = BxDFType(flags & ~BxDF_REFLECTION);
        Color f(0, 0, 0);
        for (int i = 0; i < mBxDFsSize; ++i) {
            if (bxdfs[i]->matchesFlags(flags))
                f += bxdfs[i]->f(woL, wiL);
        }
        return f;
    }
    
    virtual Color sample_f(const Direction& woW, Direction* wiW, const Sample& sample, float* pdf, BxDFType flags = BxDF_ALL, BxDFType *sampledType = nullptr) const {
        int matching_comps = getComponentsCount(flags);
        if (matching_comps == 0) {
            *pdf = 0.f;
            if (sampledType)
                *sampledType = BxDFType(0);
            return Color(0, 0, 0);
        }
        
        int which = std::min(floor2Int(sample.x * matching_comps), matching_comps-1);
        BxDF *bxdf = nullptr;
        int count = which;
        for (int i = 0; i < mBxDFsSize; ++i) {
            if (bxdfs[i]->matchesFlags(flags) && count-- == 0) {
                bxdf = bxdfs[i];
                break;
            }
        }
        Assert(bxdf != nullptr);
        
        Direction woL = world2Local(woW);
        Direction wiL;
        *pdf = 0.f;
        Color f = bxdf->sample_f(woL, &wiL, sample, pdf);
        if (*pdf == 0.f) {
            if (sampledType)
                *sampledType = BxDFType(0);
            return Color(0, 0, 0);
        }
        if (sampledType)
            *sampledType = bxdf->type;
        *wiW = local2World(wiL);
        
        // Compute overall PDF with all matching _BxDF_s
        if (!(bxdf->type & BxDF_SPECULAR) && matching_comps > 1) {
            for (int i = 0; i < mBxDFsSize; ++i) {
                if (bxdfs[i] != bxdf && bxdfs[i]->matchesFlags(flags))
                    *pdf += bxdfs[i]->pdf(woL, wiL);
            }
        }
        if (matching_comps > 1) *pdf /= matching_comps;
        
        // Compute value of BSDF for sampled direction
        if (!(bxdf->type & BxDF_SPECULAR)) {
            f = Color(0, 0, 0);
            if (Dot(*wiW, mGeometryNormal) * Dot(woW, mGeometryNormal) > 0) // ignore BTDFs
                flags = BxDFType(flags & ~BxDF_TRANSMISSION);
            else // ignore BRDFs
                flags = BxDFType(flags & ~BxDF_REFLECTION);
            for (int i = 0; i < mBxDFsSize; ++i)
                if (bxdfs[i]->matchesFlags(flags)) {
                    f += bxdfs[i]->f(woL, wiL);
                }
        }
        
        return f;
    }
    
//    // hemispherical-directional reflectance
//    virtual Color rho(const Direction& w, const std::vector<Sample>& samples) const {
//
//    }
//    
//    // hemispherical-hemispherical reflectance
//    virtual Color rho(const std::vector<Sample>& samples_o, const std::vector<Sample>& samples_i) const {
//
//    }
    
    virtual float pdf(const Direction& woW, const Direction& wiW, BxDFType flags = BxDF_ALL) const {
        if (mBxDFsSize == 0)
            return 0.;
        Direction woL = world2Local(woW);
        Direction wiL = world2Local(wiW);
        float pdf = 0.f;
        int matching_comps = 0;
        for (int i = 0; i < mBxDFsSize; ++i) {
            if (bxdfs[i]->matchesFlags(flags)) {
                ++matching_comps;
                pdf += bxdfs[i]->pdf(woL, wiL);
            }
        }
        float v = matching_comps > 0 ? pdf / matching_comps : 0.f;
        return v;
    }
};


#endif /* BSDF_hpp */
